#!/usr/bin/env node

import * as fs from "fs/promises"
import * as path from "path"
import * as grpc from "@grpc/grpc-js"
import * as protoLoader from "@grpc/proto-loader"
import chalk from "chalk"

const IMPL_FILE = path.resolve("src/generated/standalone/host-bridge-clients.ts")
const INTERFACE_FILE = path.resolve("src/generated/hosts/host-bridge-client-types.ts")
const DESCRIPTOR_SET = path.resolve("dist-standalone/proto/descriptor_set.pb")

const typeNameToFQN = new Map()

function addTypeNameToFqn(name, fqn) {
	if (typeNameToFQN.has(name)) {
		throw new Error(`Proto type ${name} redefined (${fqn}).`)
	}
	typeNameToFQN.set(name, fqn)
}
function getFqn(name) {
	if (!typeNameToFQN.has(name)) {
		throw Error(`No FQN for ${name}`)
	}
	return typeNameToFQN.get(name)
}
/**
 * Main function to generate the host bridge client
 */
async function main() {
	// Load service definitions from descriptor set
	const descriptorBuffer = await fs.readFile(DESCRIPTOR_SET)
	const packageDefinition = protoLoader.loadFileDescriptorSetFromBuffer(descriptorBuffer)
	const proto = grpc.loadPackageDefinition(packageDefinition)

	// Extract host services and proto messages from the proto definition
	const hostServices = {}
	for (const [name, def] of Object.entries(proto.host)) {
		if (def && "service" in def) {
			hostServices[name] = def
		} else {
			addTypeNameToFqn(name, `proto.host.${name}`)
		}
	}
	for (const [name, def] of Object.entries(proto.cline)) {
		if (def && !("service" in def)) {
			addTypeNameToFqn(name, `proto.cline.${name}`)
		}
	}

	// Generate interfaces file
	await generateInterfacesFile(hostServices)

	// // Generate implementation file
	await generateImplementationFile(hostServices)

	console.log(`Generated host bridge client files at:`)
	console.log(`- ${INTERFACE_FILE}`)
	console.log(`- ${IMPL_FILE}`)
}

/**
 * Generate the client interfaces file.
 */
async function generateInterfacesFile(hostServices) {
	const clientInterfaces = []
	for (const [name, def] of Object.entries(hostServices)) {
		const clientInterface = generateClientInterface(name, def)
		clientInterfaces.push(clientInterface)
	}
	const content = `// GENERATED CODE -- DO NOT EDIT!
// Generated by scripts/generate-host-bridge-client.mjs
import * as proto from "@shared/proto/index"
import { StreamingCallbacks } from "@hosts/host-provider-types"

${clientInterfaces.join("\n\n")}
`
	// Write output file
	await fs.mkdir(path.dirname(INTERFACE_FILE), { recursive: true })
	await fs.writeFile(INTERFACE_FILE, content)
}

/**
 * Generate a client interface for a service.
 */
function generateClientInterface(serviceName, serviceDefinition) {
	// Get the methods from the service definition
	const methods = Object.entries(serviceDefinition.service)
		.map(([methodName, methodDef]) => {
			const requestType = getFqn(methodDef.requestType.type.name)
			const responseType = getFqn(methodDef.responseType.type.name)

			if (!methodDef.responseStream) {
				// Generate unary method signature.
				return `	${methodName}(request: ${requestType}): Promise<${responseType}>;`
			}
			// Generate streaming method signature.
			return `	${methodName}(request: ${requestType}, callbacks: StreamingCallbacks<${responseType}>): () => void;`
		})
		.join("\n\n")

	// Generate the interface
	return `/**
 * Interface for ${serviceName} client.
 */
export interface ${serviceName}ClientInterface {

${methods}
}`
}

/**
 * Generate the client implementations file.
 */
async function generateImplementationFile(hostServices) {
	// Generate imports
	const imports = []
	// Add imports for the interfaces
	for (const [name, _def] of Object.entries(hostServices)) {
		imports.push(`import { ${name}ClientInterface } from "@generated/hosts/host-bridge-client-types"`)
	}
	const clientImplementations = []
	for (const [name, def] of Object.entries(hostServices)) {
		clientImplementations.push(generateClientImplementation(name, def))
	}

	const content = `// GENERATED CODE -- DO NOT EDIT!
// Generated by scripts/generate-host-bridge-client.mjs
import { asyncIteratorToCallbacks } from "@/standalone/utils"
import * as niceGrpc from "@generated/nice-grpc/index"
import { StreamingCallbacks } from "@hosts/host-provider-types"
import * as proto from "@shared/proto/index"
import { Channel, createClient } from "nice-grpc"

${imports.join("\n")}

${clientImplementations.join("\n\n")}
`

	// Write output file
	await fs.mkdir(path.dirname(IMPL_FILE), { recursive: true })
	await fs.writeFile(IMPL_FILE, content)
}

/**
 * Generate a client implementation class for a service
 */
function generateClientImplementation(serviceName, serviceDefinition) {
	// Get the methods from the service definition
	const methods = Object.entries(serviceDefinition.service)
		.map(([methodName, methodDef]) => {
			// Get fully qualified type names
			const requestType = getFqn(methodDef.requestType.type.name)
			const responseType = getFqn(methodDef.responseType.type.name)
			const isStreamingResponse = methodDef.responseStream

			if (!isStreamingResponse) {
				return `  ${methodName}(request: ${requestType}): Promise<${responseType}> {
    return this.client.${methodName}(request)
  }`
			} else {
				// Generate streaming method
				return `  ${methodName}(request: ${requestType}, callbacks: StreamingCallbacks<${responseType}>): () => void {
	const abortController = new AbortController()
	const stream: AsyncIterable<${responseType}> = this.client.${methodName}(request, {signal: abortController.signal})
    asyncIteratorToCallbacks(stream, callbacks)
	return () => {abortController.abort()}
  }`
			}
		})
		.join("\n\n")

	// Generate the class
	return `/**
 * Type-safe client implementation for ${serviceName}.
 */
export class ${serviceName}ClientImpl implements ${serviceName}ClientInterface {
  private client: niceGrpc.host.${serviceName}Client 

  constructor(channel: Channel) {
    this.client = createClient(niceGrpc.host.${serviceName}Definition, channel)
  }

${methods}
}`
}

// Run the main function
main().catch((error) => {
	console.error(chalk.red("Error:"), error)
	process.exit(1)
})
